#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：gas_pipeline_qc_partition_ga.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/23 14:37
gas_pipeline_qc_partition_ga
再加一条，在分析的时候是可以分开分析的，比如有瓶气含有四个组分ABCD
可以先分析A 然后BCD可以一起分析，因为可能A的时间长，BCD的时间短，
加上收敛曲线图 所有图的字用黑色
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Gas Mixing -> QC Scheduling with Component-wise Charging and Grouped Testing
----------------------------------------------------------------------------
- MIX: same recipe bottles can be batched; charging is done component-by-component
       (even for a batch) with sequence-dependent setups.
- QC : each device measures exactly ONE component. For each bottle we decide a
       grouping policy over its components:
         * ALL_TOGETHER: one 1-pass group over all components;
         * SPLIT_LONGEST: longest component alone, the rest as one 1-pass group;
         * EACH_SEPARATE: each component alone (serial).
       1-pass group occupies all its devices with the SAME [start,end] and
       duration = max(single-device times in that group). No mid-detach allowed.
- A tiny GA searches per-bottle grouping policies and plots a convergence curve.

Author: you
"""

import random
import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Tuple, Optional

try:
    import matplotlib.pyplot as plt
    HAS_MPL = True
except Exception:
    HAS_MPL = False


# ============================ CONFIG ============================ #

CONFIG = {
    # ---- scale ----
    "NUM_BOTTLES": 20,
    "NUM_MIX_MACHINES": 2,
    "NUM_QC_DEVICES": 6,
    "MAX_COMPONENTS_PER_BOTTLE": 4,

    # ---- MIX (charging) ----
    "MIX_BATCH_CAP": 300,                 # max bottles per batch (same recipe)
    "MIX_RATE_PER_H": [180.0, 150.0],     # bottles/hour per mixer (not used if per-bottle times used)
    "MIX_SETUP_SAME": 0.05,               # setup if last_component == current_component
    "MIX_SETUP_DIFF": 0.35,               # setup otherwise
    # per-bottle charging time (hours) for each component
    "COMPONENT_FILL_TIME": {
        "O2": 0.15, "CO2": 0.20, "CH4": 0.22, "CO": 0.18,
        "H2": 0.12, "N2": 0.10, "C2H6": 0.25
    },
    # charging component order priority (left is earlier); others go by alpha
    "CHARGE_PRIORITY": ["CO2", "O2", "CH4", "H2", "CO", "N2", "C2H6"],

    # ---- QC (testing) ----
    "COMPONENT_TYPES": ["O2", "CO2", "CH4", "CO", "H2", "N2", "C2H6"],
    "COMPONENT_TEST_TIME": {              # single-device time per bottle
        "O2": 0.18, "CO2": 0.22, "CH4": 0.25, "CO": 0.20,
        "H2": 0.15, "N2": 0.12, "C2H6": 0.28
    },
    "QC_SETUP_TIME": 0.12,                # change of method signature (component)

    # ---- GA over QC grouping policy ----
    "GA_POP": 10,
    "GA_GENS": 500,
    "GA_PC": 0.9,
    "GA_PM": 0.35,
    "GA_ELITE": 10,

    # ---- plotting ----
    "LABEL_WRAP_CHARS": 22,               # wrap text width; all text in black
    "RANDOM_SEED": 7
}


# ============================ Data Models ============================ #

@dataclass(frozen=True)
class Recipe:
    components: Tuple[Tuple[str, float], ...]  # (component, concentration%)

    def key(self) -> Tuple[Tuple[str,float], ...]:
        return tuple(sorted(self.components, key=lambda x: x[0]))


@dataclass
class Bottle:
    id: int
    recipe: Recipe


@dataclass
class MixTask:
    start: float
    end: float
    machine: int
    label: str


@dataclass
class QCTask:
    start: float
    end: float
    device: int
    label: str
    group_key: Optional[Tuple[int,float,float,int]] = None  # (bid, st, en, group_index) for 1-pass coloring


# ============================ Random Instance ============================ #

class RandomData:
    def __init__(self, cfg):
        self.cfg = cfg
        random.seed(cfg["RANDOM_SEED"])
        np.random.seed(cfg["RANDOM_SEED"])

    def gen(self):
        N = self.cfg["NUM_BOTTLES"]
        D = self.cfg["NUM_QC_DEVICES"]
        types = self.cfg["COMPONENT_TYPES"]
        Kmax = self.cfg["MAX_COMPONENTS_PER_BOTTLE"]

        # map device -> a single component type (cycled)
        dev_to_comp = [types[i % len(types)] for i in range(D)]

        uniq_types = list(dict.fromkeys(dev_to_comp))
        bottles: List[Bottle] = []
        for i in range(N):
            k = np.random.randint(1, Kmax+1)
            comps = random.sample(uniq_types, k)
            comps_t = tuple((c, round(np.random.uniform(0.1, 5.0), 3)) for c in comps)
            bottles.append(Bottle(i, Recipe(comps_t)))
        return bottles, dev_to_comp


# ============================ Scheduler ============================ #

class Scheduler:
    """
    - MIX: component-by-component for each batch; sequence-dependent setups by last component.
    - QC : per-bottle, we execute an ordered list of groups; each group is either 1-pass (size>=2)
           or single-component (serial-like).
    """
    def __init__(self, cfg, bottles: List[Bottle], dev_to_comp: List[str]):
        self.cfg = cfg
        self.bottles = bottles
        self.N = len(bottles)
        self.dev_to_comp = dev_to_comp
        self.D_mix = cfg["NUM_MIX_MACHINES"]
        self.D_qc = cfg["NUM_QC_DEVICES"]

    # ---------- charging component order ----------
    def charging_sequence(self, comps: List[str]) -> List[str]:
        pri = {c: idx for idx, c in enumerate(self.cfg["CHARGE_PRIORITY"])}
        return sorted(comps, key=lambda c: (pri.get(c, 10_000), c))

    # ---------- MIX: component-wise charging with batching ----------
    def schedule_mixing(self) -> Tuple[Dict[int,float], List[MixTask]]:
        cap = self.cfg["MIX_BATCH_CAP"]
        setup_same = self.cfg["MIX_SETUP_SAME"]
        setup_diff = self.cfg["MIX_SETUP_DIFF"]
        fill_t = self.cfg["COMPONENT_FILL_TIME"]

        # group by recipe key
        key2ids: Dict[Tuple, List[int]] = {}
        for b in self.bottles:
            key2ids.setdefault(b.recipe.key(), []).append(b.id)

        # split to batches
        batches: List[List[int]] = []
        for _, ids in key2ids.items():
            cur=[]
            for bid in ids:
                cur.append(bid)
                if len(cur) >= cap:
                    batches.append(cur); cur=[]
            if cur: batches.append(cur)

        # sort by size (big first)
        batches.sort(key=lambda bl: -len(bl))

        ready = [0.0]*self.cfg["NUM_MIX_MACHINES"]
        last_comp = [None]*self.cfg["NUM_MIX_MACHINES"]
        mix_tasks: List[MixTask] = []
        mix_end: Dict[int,float] = {}

        for bl in batches:
            comps = [c for c,_ in self.bottles[bl[0]].recipe.components]
            seq = self.charging_sequence(comps)
            batch_prev_end = 0.0
            for comp in seq:
                t_per_bottle = float(fill_t.get(comp, 0.02))
                pt = t_per_bottle * len(bl)

                best=None
                for m in range(self.cfg["NUM_MIX_MACHINES"]):
                    stp = setup_same if (last_comp[m] == comp and last_comp[m] is not None) else \
                          (0.0 if last_comp[m] is None else setup_diff)
                    st = max(ready[m] + stp, batch_prev_end)
                    en = st + pt
                    if (best is None) or (en < best[0]):
                        best=(en,m,st,pt,stp)
                en,m,st,pt_use,stp = best
                mix_tasks.append(MixTask(st, en, m, f"Batch({','.join('S'+str(i) for i in bl)})-{comp}"))
                ready[m] = en; last_comp[m] = comp
                batch_prev_end = en

            for bid in bl:
                mix_end[bid] = batch_prev_end

        return mix_end, mix_tasks

    # ---------- QC single times ----------
    def single_test_time(self, comp: str) -> float:
        return float(self.cfg["COMPONENT_TEST_TIME"].get(comp, 0.2))

    # ---------- device candidates per component ----------
    def comp_to_devs(self) -> Dict[str, List[int]]:
        mapping: Dict[str, List[int]] = {}
        for d, c in enumerate(self.dev_to_comp):
            mapping.setdefault(c, []).append(d)
        return mapping

    # ---------- grouping policies per bottle ----------
    # policy codes:
    #   0 = ALL_TOGETHER        -> [ {all components} ]
    #   1 = SPLIT_LONGEST       -> [ {longest}, {rest} ]
    #   2 = EACH_SEPARATE       -> [ {c1}, {c2}, ... ] (order by device id)
    def make_groups(self, comps: List[str], policy_code: int) -> List[List[str]]:
        if len(comps) <= 1:
            return [comps[:]]
        times = {c: self.single_test_time(c) for c in comps}
        if policy_code == 0:  # ALL together
            return [comps[:]]
        if policy_code == 1:  # SPLIT longest
            longest = max(comps, key=lambda c: times[c])
            rest = [c for c in comps if c != longest]
            return [[longest], rest] if rest else [[longest]]
        # EACH separate
        # stable by device id if possible
        c2d = self.comp_to_devs()
        comps_sorted = sorted(comps, key=lambda c: (min(c2d.get(c,[999])) if c2d.get(c) else 999, c))
        return [[c] for c in comps_sorted]

    # ---------- QC scheduler given grouping list per bottle ----------
    def schedule_qc_with_groups(self, mix_end: Dict[int,float], bottle_groups: Dict[int, List[List[str]]]
                               ) -> Tuple[float, List[QCTask]]:
        D = self.D_qc
        ready = [0.0]*D
        last_sig = [None]*D
        setup_t = self.cfg["QC_SETUP_TIME"]
        comp2devs = self.comp_to_devs()

        tasks: List[QCTask] = []
        # bottles in order of mix_end
        order_seq = sorted(self.bottles, key=lambda b: mix_end[b.id])

        for b in order_seq:
            bid = b.id
            t_cur = mix_end[bid]
            groups = bottle_groups[bid]  # ordered list of component lists

            for gi, group in enumerate(groups):
                # pick device for each component in this group (earliest ready+setup)
                chosen = []
                for comp in group:
                    cands = comp2devs.get(comp, [])
                    if not cands:
                        raise RuntimeError(f"No QC device for component {comp}")
                    best_d=None; best_ready=None
                    for d in cands:
                        stp = 0.0 if (last_sig[d] is None or last_sig[d]==comp) else setup_t
                        rdy = ready[d] + stp
                        if (best_ready is None) or (rdy < best_ready):
                            best_ready = rdy; best_d=d
                    chosen.append((best_d, comp))

                if len(group) >= 2:
                    # 1-pass group: same start & end across all chosen devices
                    earliest = t_cur
                    per=[]
                    for d, comp in chosen:
                        stp = 0.0 if (last_sig[d] is None or last_sig[d]==comp) else setup_t
                        earliest = max(earliest, ready[d] + stp)
                        pt = self.single_test_time(comp)
                        per.append((d, comp, stp, pt))
                    seg_len = max(pt for (_,_,_,pt) in per)
                    st = earliest; en = st + seg_len
                    group_key = (bid, round(st,6), round(en,6), gi)
                    used = "+".join(f"E{d+1}" for (d,_,_,_) in per)
                    label = f"S{bid} 1-pass[{used}] max={seg_len:.2f}h"
                    for d, comp, stp, pt in per:
                        tasks.append(QCTask(st, en, d, label, group_key))
                        ready[d] = en; last_sig[d] = comp
                    t_cur = en
                else:
                    # single component: serial-like
                    d, comp = chosen[0]
                    stp = 0.0 if (last_sig[d] is None or last_sig[d]==comp) else setup_t
                    st = max(t_cur, ready[d] + stp)
                    pt = self.single_test_time(comp)
                    en = st + pt
                    tasks.append(QCTask(st, en, d, f"S{bid} serial E{d+1}({comp})", None))
                    ready[d] = en; last_sig[d] = comp
                    t_cur = en

        Cmax = max((t.end for t in tasks), default=0.0)
        return Cmax, tasks


# ============================ GA for grouping policy ============================ #

class GA:
    """
    Gene per bottle: policy code in {0,1,2}
      0=ALL_TOGETHER, 1=SPLIT_LONGEST, 2=EACH_SEPARATE
    Fitness: minimize Cmax (you can extend with tardiness etc.)
    """
    def __init__(self, cfg, scheduler: Scheduler, mix_end: Dict[int,float]):
        self.cfg = cfg
        self.S = scheduler
        self.mix_end = mix_end
        self.N = self.S.N
        random.seed(cfg["RANDOM_SEED"])
        np.random.seed(cfg["RANDOM_SEED"])

    def rand_ind(self):
        return np.random.randint(0,3,size=self.N,dtype=int)  # 0/1/2 per bottle

    def crossover(self, a, b):
        n=len(a); cut=np.random.randint(1,n)
        c1=np.concatenate([a[:cut], b[cut:]])
        c2=np.concatenate([b[:cut], a[cut:]])
        return c1,c2

    def mutate(self, ind, pm=0.2):
        if np.random.rand()<pm:
            i=np.random.randint(self.N)
            ind[i] = np.random.randint(0,3)
        return ind

    def ind_to_groups(self, ind) -> Dict[int, List[List[str]]]:
        mapping={}
        for i, b in enumerate(self.S.bottles):
            comps=[c for c,_ in b.recipe.components]
            mapping[b.id]=self.S.make_groups(comps, int(ind[i]))
        return mapping

    def fitness(self, ind):
        groups = self.ind_to_groups(ind)
        Cmax, _ = self.S.schedule_qc_with_groups(self.mix_end, groups)
        return Cmax

    def run(self):
        pop=self.cfg["GA_POP"]; gens=self.cfg["GA_GENS"]
        pc=self.cfg["GA_PC"]; pm=self.cfg["GA_PM"]; elite=self.cfg["GA_ELITE"]
        # init
        pool=[]
        for _ in range(pop):
            ind=self.rand_ind()
            f=self.fitness(ind)
            pool.append((ind,f))
        best=min(pool, key=lambda x:x[1])
        hist=[best[1]]

        for g in range(1, gens+1):
            pool.sort(key=lambda x:x[1])
            elites=pool[:elite]
            new=[(e[0].copy(), e[1]) for e in elites]
            while len(new)<pop:
                p1=random.choice(pool)[0]
                p2=random.choice(pool)[0]
                c1,c2=p1.copy(), p2.copy()
                if np.random.rand()<pc:
                    c1,c2=self.crossover(p1,p2)
                if np.random.rand()<pm:
                    c1=self.mutate(c1)
                if np.random.rand()<pm:
                    c2=self.mutate(c2)
                new.append((c1, self.fitness(c1)))
                if len(new)<pop:
                    new.append((c2, self.fitness(c2)))
            pool=new
            cur=min(pool, key=lambda x:x[1])
            if cur[1] < best[1]:
                best=cur
            hist.append(best[1])
        return best[0], best[1], hist


# ============================ Plotting ============================ #

def wrap_text(s: str, width: int) -> str:
    # simple char-based wrap; break at spaces when available
    out=[]; line=""
    for ch in s:
        line += ch
        if len(line) >= width and ch == ' ':
            out.append(line.rstrip()); line=""
    if line: out.append(line)
    return "\n".join(out)

def plot_gantt(mix_tasks: List[MixTask], qc_tasks: List[QCTask], cfg,
               title="Gantt: MIX (component-wise) + QC (grouped 1-pass/serial)"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip plotting)"); return

    # rows
    rows=[]
    for m in range(cfg["NUM_MIX_MACHINES"]):
        rows.append(("MIX", m))
    for d in range(cfg["NUM_QC_DEVICES"]):
        rows.append(("QC", d))
    row_idx={("MIX",m): m for m in range(cfg["NUM_MIX_MACHINES"])}
    off=cfg["NUM_MIX_MACHINES"]
    for d in range(cfg["NUM_QC_DEVICES"]):
        row_idx[("QC",d)] = off + d

    # colors for 1-pass groups
    from itertools import cycle
    palette = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
    grp_color: Dict[Tuple[int,float,float,int], str] = {}
    for t in qc_tasks:
        if t.group_key and t.group_key not in grp_color:
            grp_color[t.group_key] = next(palette)

    fig_h = max(4.0, 0.6*len(rows)+2)
    fig, ax = plt.subplots(figsize=(12, fig_h))

    # MIX
    for t in mix_tasks:
        y=row_idx[("MIX",t.machine)]
        ax.barh(y, t.end-t.start, left=t.start, height=0.8,
                edgecolor='black', linewidth=0.6)
        ax.text((t.start+t.end)/2, y, wrap_text(t.label, cfg["LABEL_WRAP_CHARS"]),
                ha='center', va='center', color='black', fontsize=8)

    # QC
    for t in qc_tasks:
        y=row_idx[("QC",t.device)]
        color = grp_color.get(t.group_key, None)
        ax.barh(y, t.end-t.start, left=t.start, height=0.8,
                color=color, edgecolor='black', linewidth=0.6)
        ax.text((t.start+t.end)/2, y, wrap_text(t.label, cfg["LABEL_WRAP_CHARS"]),
                ha='center', va='center', color='black', fontsize=8)

    # y labels
    yticks=[]; ylabels=[]
    for m in range(cfg["NUM_MIX_MACHINES"]):
        yticks.append(row_idx[("MIX", m)])
        ylabels.append(f"MIX-M{m+1}")
    types = CONFIG["COMPONENT_TYPES"]
    for d in range(cfg["NUM_QC_DEVICES"]):
        yticks.append(row_idx[("QC", d)])
        ylabels.append(f"QC-E{d+1} ({types[d % len(types)]})")

    ax.set_yticks(yticks); ax.set_yticklabels(ylabels, color='black')
    ax.set_xlabel("Time (h)", color='black'); ax.set_title(title, color='black')
    ax.tick_params(axis='x', colors='black'); ax.tick_params(axis='y', colors='black')
    ax.grid(axis='x', alpha=0.3, color='black')
    for spine in ax.spines.values():
        spine.set_edgecolor('black')
    plt.tight_layout(); plt.show()

def plot_convergence(hist, title="GA Convergence"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip plotting)"); return
    plt.figure(figsize=(6,4))
    plt.plot(hist, linewidth=2)
    plt.xlabel("Generation", color='black'); plt.ylabel("Best Cmax", color='black')
    plt.title(title, color='black'); plt.grid(True, alpha=0.3, color='black')
    ax=plt.gca()
    ax.tick_params(axis='x', colors='black'); ax.tick_params(axis='y', colors='black')
    for spine in ax.spines.values():
        spine.set_edgecolor('black')
    plt.tight_layout(); plt.show()


# ============================ Main ============================ #

def main():
    random.seed(CONFIG["RANDOM_SEED"]); np.random.seed(CONFIG["RANDOM_SEED"])

    # 1) instance
    bottles, dev_to_comp = RandomData(CONFIG).gen()

    # 2) MIX schedule (component-wise)
    sch = Scheduler(CONFIG, bottles, dev_to_comp)
    mix_end, mix_tasks = sch.schedule_mixing()

    # 3) GA over QC grouping policies
    ga = GA(CONFIG, sch, mix_end)
    best_ind, best_cmax, hist = ga.run()

    # 4) Decode best to concrete QC schedule and plot
    #    (build groups per bottle from best_ind)
    bottle_groups={}
    for i,b in enumerate(bottles):
        comps=[c for c,_ in b.recipe.components]
        bottle_groups[b.id] = sch.make_groups(comps, int(best_ind[i]))

    Cmax, qc_tasks = sch.schedule_qc_with_groups(mix_end, bottle_groups)

    print(f"Bottles: {CONFIG['NUM_BOTTLES']}, Mixers: {CONFIG['NUM_MIX_MACHINES']}, QC devices: {CONFIG['NUM_QC_DEVICES']}")
    print("Best grouping policy codes per bottle (0=ALL,1=SPLIT_LONGEST,2=EACH):")
    print(best_ind.tolist())
    print(f"Best Cmax: {Cmax:.3f} h")

    if HAS_MPL:
        plot_convergence(hist, title="GA Convergence (minimize Cmax via grouping policies)")
        plot_gantt(mix_tasks, qc_tasks, CONFIG)
    else:
        print("(Install matplotlib to view plots.)")

if __name__ == "__main__":
    main()
